import argparse
import os
import cv2
import numpy as np
import torch


def str2bool(v):
    if v.lower() in ['true', 1]:
        return True
    elif v.lower() in ['false', 0]:
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def count_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def compute_class_weights(mask_dir, num_classes, exponent=0.5):
    """
    计算类别权重，用于处理类别不平衡
    Args:
        mask_dir: 掩码文件目录
        num_classes: 类别数量
        exponent: 权重指数（1.0表示平衡权重，<1.0增加罕见类别的权重）
    Returns:
        class_weights: 每个类别的权重 (num_classes,)
    """
    pixel_counts = np.zeros(num_classes)

    for mask_file in os.listdir(mask_dir):
        mask_path = os.path.join(mask_dir, mask_file)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        for class_idx in range(num_classes):
            pixel_counts[class_idx] += np.sum(mask == class_idx)

    total_pixels = np.sum(pixel_counts)
    class_frequencies = pixel_counts / total_pixels
    class_weights = 1.0 / (class_frequencies ** exponent + 1e-8)
    class_weights = class_weights / np.sum(class_weights) * num_classes

    return torch.tensor(class_weights, dtype=torch.float32)


def visualize_predictions(image, prediction, target=None, num_classes=7):
    """
    可视化预测结果
    Args:
        image: 原始图像 (H, W, 3)
        prediction: 预测分割图 (H, W)
        target: 真实分割图 (H, W) 可选
        num_classes: 类别数量
    Returns:
        vis_image: 可视化图像
    """
    # 类别颜色映射
    color_map = {
        0: [0, 0, 0],  # 背景 - 黑色
        1: [255, 0, 0],  # 建筑 - 红色
        2: [0, 255, 0],  # 道路 - 绿色
        3: [0, 0, 255],  # 水域 - 蓝色
        4: [255, 255, 0],  # 荒地 - 黄色
        5: [0, 255, 255],  # 森林 - 青色
        6: [255, 0, 255]  # 农田 - 洋红色
    }

    # 创建彩色预测图
    h, w = prediction.shape
    color_pred = np.zeros((h, w, 3), dtype=np.uint8)
    for class_idx, color in color_map.items():
        if class_idx < num_classes:
            color_pred[prediction == class_idx] = color

    # 创建彩色真实图（如果有）
    if target is not None:
        color_target = np.zeros((h, w, 3), dtype=np.uint8)
        for class_idx, color in color_map.items():
            if class_idx < num_classes:
                color_target[target == class_idx] = color

    # 组合可视化
    if target is not None:
        # 三列：原始图像、真实标注、预测结果
        vis_image = np.concatenate([image, color_target, color_pred], axis=1)
    else:
        # 两列：原始图像、预测结果
        vis_image = np.concatenate([image, color_pred], axis=1)

    return vis_image


def save_prediction(image, prediction, output_path, num_classes=7):
    """
    保存预测结果图像
    Args:
        image: 原始图像 (H, W, 3)
        prediction: 预测分割图 (H, W)
        output_path: 输出文件路径
        num_classes: 类别数量
    """
    vis_image = visualize_predictions(image, prediction, num_classes=num_classes)
    cv2.imwrite(output_path, cv2.cvtColor(vis_image, cv2.COLOR_RGB2BGR))


def create_comparison_figure(image, prediction, target, output_path, num_classes=7):
    """
    创建预测与真实标注的对比图
    Args:
        image: 原始图像 (H, W, 3)
        prediction: 预测分割图 (H, W)
        target: 真实分割图 (H, W)
        output_path: 输出文件路径
        num_classes: 类别数量
    """
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    # 原始图像
    axes[0].imshow(image)
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    # 真实标注
    color_target = np.zeros_like(image)
    for class_idx in range(num_classes):
        mask = target == class_idx
        color_target[mask] = visualize_predictions(None, target, num_classes=num_classes)[mask]
    axes[1].imshow(color_target)
    axes[1].set_title('Ground Truth')
    axes[1].axis('off')

    # 预测结果
    color_pred = np.zeros_like(image)
    for class_idx in range(num_classes):
        mask = prediction == class_idx
        color_pred[mask] = visualize_predictions(None, prediction, num_classes=num_classes)[mask]
    axes[2].imshow(color_pred)
    axes[2].set_title('Prediction')
    axes[2].axis('off')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()